#!/usr/bin/env python

from Ganga.GPIDev.Adapters.IBackend import IBackend
from Ganga.GPIDev.Schema import *
from Ganga.GPIDev.Lib.File import FileBuffer

#### Additional external libraries
import xml.etree.ElementTree as ET
from CloudDrone.CloudServer import CloudServer
from CloudDrone.VM import VM
from CloudDrone.Template import Template 

#### Adding plugin configuration variables 
import Ganga.Utility.Config
config = Ganga.Utility.Config.makeConfig('OpenNebula','parameters of the OpenNebula backend')
config.addOption( 'user', 'hurng-chun.lee@cern.ch', 'OpenNebula user id')

#### Retriving logging facility
import Ganga.Utility.logging
logger = Ganga.Utility.logging.getLogger()

class CloudServerAdv(CloudServer):
    """
    Extends CloudDrone.CloudServer with more VM and template management methods
    """
    def __init__(self):
        super(CloudServerAdv, self).__init__()
        self.user   = None
        self.ugroup = None
        self.__connect__()

    def allocate_template(self, str_template):
        rslt = self.one.template.allocate(str_template)
        return rslt

    def delete_template(self, template_id):
        rslt = self.one.template.delete(template_id)
        return rslt

    def __connect__(self):
        self.set_credentials(config['user'],"password")

        ## get user info and test if the credential information is correct
        info_xml = self.one.user.info(-1)
        if info_xml.find('ID') ==  None:
            logger.error('cannot authorize user: %s' % config['user'] )
        else:
            self.user   = config['user']
            self.ugroup = info_xml.find('GNAME').text
            logger.debug("user: %s group: %s" % ( self.user, self.ugroup ))

        ## remove account information
        del info_xml

class OpenNebula(IBackend):
    """
    Start VM images on the OpenNebula cloud.
    """
    _schema = Schema(Version(0,1), {'id'           : SimpleItem(defvalue=-1,protected=1,copyable=0,doc='VM instance id'),
                                    'status'       : SimpleItem(defvalue='',protected=1,copyable=0,doc='VM instance status'),
                                    'state_nr'     : SimpleItem(defvalue=-1,protected=1,copyable=0,hidden=1,doc='VM state number'),
                                    'lcm_state_nr' : SimpleItem(defvalue=-1,protected=1,copyable=0,hidden=1,doc='VM lcm state number'),
                                    'template_id'  : SimpleItem(defvalue=-1,protected=1,copyable=0,hidden=1,doc='VM template id'),
                                    'cpu'          : SimpleItem(defvalue=1 ,copyable=1,doc='number of CPU per VM instance'),
                                    'mem'          : SimpleItem(defvalue=8192,copyable=1,doc='size of memory in MB per VM instance'),
                                    'public_net'   : SimpleItem(defvalue=True,copyable=1,doc='enable the public network interface'),
                                    'exitcode'     : SimpleItem(defvalue=None,typelist=['int','type(None)'],protected=1,copyable=0,doc='Process exit code.'),
                                    'vm_host'      : SimpleItem(defvalue=[],typelist=['str'],sequence=1,strict_sequence=0,protected=1,copyable=0,doc='hostname/ip of the VM instance'),
                                    'actualCE'     : SimpleItem(defvalue='',protected=1,copyable=0,doc='String representation of backend.vm_host'),
                                    'reason'       : SimpleItem(defvalue='',protected=1,copyable=0,doc='Reason of the VM status')
                                    })

    _category = 'backends'
    _name = 'OpenNebula' 
    _exportmethods = ['list_templates','stop_vm', 'resume_vm']

    def __init__(self):
        super(OpenNebula,self).__init__()

    def __connect__(self):
        '''
        checks server connection and return it for internal usage
        '''
        c = clouds['OpenNebula']
        if not c or c.user != config['user']: 
            ## reconnect and update user information if necessary
            c.__connect__()
        return c

    def submit(self, jobconfig, master_input_sandbox):
        ick = False

        ## prepare the template file
        f_vm_template=self.preparejob(jobconfig,master_input_sandbox)
        logger.debug('VM tempalte created: %s' % f_vm_template)

        ## connect to the server and create a VM according to the template
        c = self.__connect__()

        str_template = ''
        f = open(f_vm_template,'r') 
        for l in f:
            str_template += '%s\n' % l.strip()
        f.close()

        self.template_id = self.create_template(str_template)
        vm   = self.create_vm( self.template_id )

        if vm:

            vm_histories = {}
            num_events   = 0;
            for h in list( vm.xml_find('HISTORY_RECORDS') ):
                num_events += 1
                seq    = int( h.findtext('SEQ') )
                host   = h.findtext('HOSTNAME')
                reason = h.findtext('REASON')
                vm_histories[seq] = {'host':host, 'reason':reason}

            ick               = True
            self.id           = vm.id_nr
            self.vm_host      = vm.ips
            self.status       = vm.state_string
            self.state_nr     = vm.state_nr
            self.lcm_state_nr = vm.lcm_state_nr

            if num_events > 0:
                self.actualCE = vm_histories[num_events-1]['host']
                self.reason   = repr(vm_histories)

            OpenNebula.updateGangaJobStatus(self.getJobObject(), self.status)

        return ick 

    def __make_vm_template__(self):
        template = '''NAME=###TEMPLATE_NAME###
CPU=1
MEMORY=###M_MEM###
DISK=[ BUS=ide, IMAGE=###VM_IMG###, IMAGE_UNAME=###VM_IMG_UNAME###, TARGET=hda ]
DISK=[ BUS=ide, IMAGE=scratch, IMAGE_UNAME=###VM_IMG_UNAME###, TARGET=hdd ]
GRAPHICS=[TYPE=vnc]
OS=[ ARCH=###VM_ARCH###, BOOT=hd ]
RAW=[ TYPE=###VM_TYPE### ]
VCPU=###N_CPU###'''

        return template

    def preparejob(self,jobconfig,master_input_sandbox):

        job  = self.getJobObject()
        inpw = job.getInputWorkspace()

        c = self.__connect__()

        ## replace template variables by values given in application ##
        template = self.__make_vm_template__()
        template = template.replace( '###TEMPLATE_NAME###', '_'.join([job.application._name,job.getFQID('.')]) ) 
        template = template.replace( '###N_CPU###' , repr(self.cpu))
        template = template.replace( '###M_MEM###' , repr(self.mem))
        template = template.replace( '###VM_ARCH###', job.application.vm_arch)
        template = template.replace( '###VM_IMG###' , job.application.vm_image)
        template = template.replace( '###VM_TYPE###', job.application.vm_type)
        template = template.replace( '###VM_IMG_UNAME###', c.user)

        ## attaching networks
        if self.public_net:
            template += '\nNIC=[ MODEL=rtl8139, NETWORK=internet, NETWORK_FILTER=132, NETWORK_UNAME=oneadmin ]'
        template += '\nNIC=[ MODEL=virtio, NETWORK=%s, NETWORK_UNAME=%s ]' % (c.ugroup, c.user)

        ## special contextualization given by the application, assuming that cloud application has a "get_context()" method
        template += '\n%s' % jobconfig.getVMContext()

        return job.getInputWorkspace().writefile(FileBuffer('__vm_template__',template),executable=0)

    def kill(self):
        logger.info("shutdown VM %d ..." % self.id)
        self.stop_vm( action='shutdown', id_vm=self.id )
        self.del_template( id_template=self.template_id )
        return 1

    #### helper functions 
    def list_templates(self):
        '''
        list available templates
        '''
        c = self.__connect__()
        for t in c.iter_templates():
            print t
            print '  |-- %s' % t.xml_str()

    def get_template(self, id_template):
        '''
        retrieve template information into Template object
        '''
        c = self.__connect__()
        return c.get_template( id_template )

    def del_template(self, id_template):
        '''
        delete Template object associated with the job
        '''
        c = self.__connect__()
        return c.delete_template( id_template )

    def create_template(self, str_template):
        '''
        create new template in OpenNebula
        '''
        c = self.__connect__()
        return c.allocate_template(str_template)

    def create_vm(self, id_template, name=None):
        '''
        instantiate a VM given a template id 
        '''
        c = self.__connect__()
        id = c._one_template_instantiate(id_template, name)
        return c.get_vm( id )

    def stop_vm(self, action='shutdown', id_vm=None):
        '''
        stop a VM, action can be 'shutdown', 'cancel', 'suspend', 'stop'
        '''

        ick = False

        if action not in ['shutdown', 'cancel', 'suspend', 'stop']:
            logger.error('unsupported VM stop action: %s' % action)
            return ick

        c = self.__connect__()

        if not id_vm:
            id_vm = self.id

        vm = c.get_vm( id_vm )

        if vm.state_string.split('/')[1] != 'RUNNING':
            logger.warning('cannot %s VM in status %s' % ( action, vm.state_string) )           
        else:
            rslt = c._one_vm_action(action, id_vm)
            if rslt != id_vm:
                logger.error('%s VM failure %s' % (action, repr(rslt)))
            else:
                ick = True
        return ick 

    def resume_vm(self, id_vm=None):
        '''
        resume a VM
        '''
        c = self.__connect__()

        if not id_vm:
            id_vm = self.id

        vm = c.get_vm( id_vm )

        if vm.state_string.split('/')[0] not in ['STOPPED','SUSPENDED']:
            logger.warning('cannot resume VM in status %s' % vm.state_string) 
        else:
            rslt = c._one_vm_action('resume', id_vm)
            if rslt != id_vm:
                logger.error('resume VM failure %s' % repr(rslt) )
            else:
                ick = True
        return ick

    def updateMonitoringInformation(jobs):
        c = clouds['OpenNebula']
        if c:
            for j in jobs:
                logger.debug("monitor VM %d ..." % j.backend.id)
                vm = c.get_vm( j.backend.id )
                if vm:
                    vm.refresh()

                    vm_histories = {} 
                    num_events   = 0;      
                    for h in list( vm.xml_find('HISTORY_RECORDS') ):
                        num_events += 1
                        seq    = int( h.findtext('SEQ') )
                        host   = h.findtext('HOSTNAME')
                        reason = h.findtext('REASON')
                        vm_histories[seq] = {'host':host, 'reason':reason}

                    j.backend.vm_host      = vm.ips
                    j.backend.status       = vm.state_string
                    j.backend.reason       = repr( vm_histories )
                    j.backend.lcm_state_nr = vm.lcm_state_nr

                    if num_events > 0:
                        j.backend.actualCE     = vm_histories[num_events-1]['host']
                        j.backend.state_nr     = vm.state_nr

                    OpenNebula.updateGangaJobStatus(j, j.backend.status)
        else:
            logger.warning('OpenNebula interface not initiated. Ignore monitoring loop.')

    updateMonitoringInformation = staticmethod(updateMonitoringInformation)

    def updateGangaJobStatus(job, status):
        '''map OpenNebula VM status to Ganga job status'''

        state_info = status.split('/')

        new_status = job.status

        if state_info[0] in ['INIT', 'PENDING', 'HOLD']:
            new_status = 'submitted'

        elif state_info[0] in ['ACTIVE']:

            if state_info[1] in ['RUNNING']:
                new_status = 'running'
            #else:
            #    new_status = 'submitted'
     
        elif state_info[0] in ['DONE']:
            new_status = 'completed'
      
        elif state_info[0] in ['FAILED']:
            new_status = 'failed'
     
        elif state_info[0] in ['STOPPED','SUSPENDED']:
            new_status = 'running'
     
        else:
            logger.warning('Unexpected VM status "%s"', status)

        ## update job to the new status only if it is necessary
        if new_status != job.status:
            job.updateStatus( new_status )

    updateGangaJobStatus = staticmethod(updateGangaJobStatus)

#### startup cloud interface object
clouds = {'OpenNebula': None}
clouds['OpenNebula'] = CloudServerAdv()
